Relation to Prior Work: The authors both discuss the implementation differences with and compare the performance of their library to strong baselines in many different application areas. Their results are impressive, especially given that some of the baselines are heavily optimized for specific problems (e.g. I'm wondering if PyTorch-Geometric's main competitor DGL should be an additional comparison point for the geometric deep learning benchmarks; I think it's often faster in practice although it may be too specialized for these architectures. I would like to see more discussion of the similarities and differences between your implementation and deep learning compilers like XLA and TVM. For instance, does your package do just-in-time CUDA code generation/compilation or perform operator fusion?

Geometric methods rely on tensors that can be encoded using a symbolic formula and data arrays, such as kernel and distance matrices. We present an extension for standard machine learning frameworks that provides comprehensive support for this abstraction on CPUs and GPUs: our toolbox combines a versatile, transparent user interface with fast runtimes and low memory usage. Unlike general purpose acceleration frameworks such as XLA, our library turns generic Python code into binaries whose performances are competitive with state-of-the-art geometric libraries - such as FAISS for nearest neighbor search - with the added benefit of flexibility. We perform an extensive evaluation on a broad class of problems: Gaussian modelling, K-nearest neighbors search, geometric deep learning, non-Euclidean embeddings and optimal transport theory. In practice, for geometric problems that involve 1k to 1M samples in dimension 1 to 100, our library speeds up baseline GPU implementations by up to two orders of magnitude.